import cv2
import numpy as np
import os
import re
class MangaPanelExtractor:
    def __init__(self, white_threshold=230, min_line_thickness=3):
        self.white_threshold = white_threshold  # 降低白色区域的阈值
        self.min_line_thickness = min_line_thickness  # 降低最小分割线宽度
        self.debug = True

    def preprocess_image(self, image):
        """预处理图片"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊减少噪点
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        
        # 使用自适应阈值处理
        binary = cv2.adaptiveThreshold(
            blurred,
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,  # 使用均值而不是高斯
            cv2.THRESH_BINARY,
            15,  # 增大邻域大小
            5
        )
        
        # 创建一个较小的核用于形态学操作
        kernel = np.ones((2,2), np.uint8)
        
        # 先进行腐蚀操作，扩大黑色区域
        eroded = cv2.erode(binary, kernel, iterations=1)
        
        # 再进行膨胀操作，但使用更大的核
        kernel_dilate = np.ones((3,3), np.uint8)
        dilated = cv2.dilate(eroded, kernel_dilate, iterations=1)
        
        return dilated

    def find_white_regions(self, gray_image):
        """查找白色区域"""
        # 使用简单阈值处理
        _, binary = cv2.threshold(gray_image, self.white_threshold, 255, cv2.THRESH_BINARY)
        
        # 对二值图像进行形态学操作
        kernel = np.ones((2,2), np.uint8)
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=1)
        
        return binary

    def find_horizontal_separators(self, binary_image):
        """查找水平分割线"""
        height, width = binary_image.shape
        horizontal_lines = []
        
        # 计算每行的白色像素数量
        row_white_pixels = np.sum(binary_image == 255, axis=1)
        
        # 找到连续的白色行
        start_y = None
        min_white_ratio = 0.6  # 进一步降低白色像素的比例要求
        
        for y in range(height):
            if row_white_pixels[y] > width * min_white_ratio:  # 如果60%以上是白色
                if start_y is None:
                    start_y = y
            elif start_y is not None:
                if y - start_y >= self.min_line_thickness:
                    horizontal_lines.append((start_y, y))
                start_y = None
        
        # 如果最后还有未处理的白色区域
        if start_y is not None and height - start_y >= self.min_line_thickness:
            horizontal_lines.append((start_y, height))
        
        return horizontal_lines

    def find_vertical_separators(self, binary_image):
        """查找垂直分割线"""
        height, width = binary_image.shape
        vertical_lines = []
        
        # 计算每列的白色像素数量
        col_white_pixels = np.sum(binary_image == 255, axis=0)
        
        # 找到连续的白色列
        start_x = None
        min_white_ratio = 0.6  # 进一步降低白色像素的比例要求
        
        for x in range(width):
            if col_white_pixels[x] > height * min_white_ratio:  # 如果60%以上是白色
                if start_x is None:
                    start_x = x
            elif start_x is not None:
                if x - start_x >= self.min_line_thickness:
                    vertical_lines.append((start_x, x))
                start_x = None
        
        # 如果最后还有未处理的白色区域
        if start_x is not None and width - start_x >= self.min_line_thickness:
            vertical_lines.append((start_x, width))
        
        return vertical_lines

    def merge_close_lines(self, lines, min_gap=30):  # 增加最小间隔
        """合并距离较近的分割线"""
        if not lines:
            return lines
            
        merged = []
        current_start, current_end = lines[0]
        
        for start, end in lines[1:]:
            if start - current_end <= min_gap:
                current_end = end
            else:
                merged.append((current_start, current_end))
                current_start, current_end = start, end
        
        merged.append((current_start, current_end))
        return merged

    def extract_panels(self, image_path, output_dir):
        """提取漫画分镜"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"\n开始处理图片: {image_path}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
            
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 预处理图片
        print("预处理图片...")
        preprocessed = self.preprocess_image(image)
        
        # 查找白色区域
        binary = self.find_white_regions(preprocessed)
        
        # 查找水平和垂直分割线
        horizontal_lines = self.find_horizontal_separators(binary)
        vertical_lines = self.find_vertical_separators(binary)
        
        # 合并距离较近的分割线
        horizontal_lines = self.merge_close_lines(horizontal_lines)
        vertical_lines = self.merge_close_lines(vertical_lines)
        
        # 在调试图像上绘制分割线
        debug_image = image.copy()
        
        # 绘制水平分割线
        for y1, y2 in horizontal_lines:
            cv2.line(debug_image, (0, y1), (width, y1), (0, 255, 0), 2)
            cv2.line(debug_image, (0, y2), (width, y2), (0, 255, 0), 2)
            
        # 绘制垂直分割线
        for x1, x2 in vertical_lines:
            cv2.line(debug_image, (x1, 0), (x1, height), (0, 255, 0), 2)
            cv2.line(debug_image, (x2, 0), (x2, height), (0, 255, 0), 2)
        
        # 保存调试图像
        cv2.imwrite(os.path.join(output_dir, "detected_lines.jpg"), debug_image)
        cv2.imwrite(os.path.join(output_dir, "binary.jpg"), binary)
        
        # 提取分镜
        panels = []
        
        # 确保有至少一个水平和垂直分割线
        if not horizontal_lines:
            horizontal_lines = [(0, height)]
        if not vertical_lines:
            vertical_lines = [(0, width)]
            
        # 添加图片边界
        if horizontal_lines[0][0] > 0:
            horizontal_lines.insert(0, (0, 0))
        if horizontal_lines[-1][1] < height:
            horizontal_lines.append((height, height))
            
        if vertical_lines[0][0] > 0:
            vertical_lines.insert(0, (0, 0))
        if vertical_lines[-1][1] < width:
            vertical_lines.append((width, width))
        
        # 根据分割线提取分镜
        panel_count = 0
        for i in range(len(horizontal_lines)-1):
            for j in range(len(vertical_lines)-1):
                y1 = horizontal_lines[i][1]
                y2 = horizontal_lines[i+1][0]
                x1 = vertical_lines[j][1]
                x2 = vertical_lines[j+1][0]
                
                # 检查区域大小
                if (x2 - x1) < 30 or (y2 - y1) < 30:  # 降低最小尺寸要求
                    continue
                    
                # 提取分镜
                panel = image[y1:y2, x1:x2]
                
                # 检查是否为空白区域
                if np.mean(panel) > self.white_threshold:
                    continue
                
                # 检查区域是否包含足够的非白色像素
                gray_panel = cv2.cvtColor(panel, cv2.COLOR_BGR2GRAY)
                non_white_ratio = np.sum(gray_panel < self.white_threshold) / (gray_panel.shape[0] * gray_panel.shape[1])
                if non_white_ratio < 0.1:  # 如果非白色像素太少，跳过
                    continue
                
                panel_count += 1
                output_path = os.path.join(output_dir, f"panel_{panel_count}.jpg")
                cv2.imwrite(output_path, panel)
                print(f"保存分镜 {panel_count}: {output_path}")
                
        return panel_count

def natural_sort_key(s):
    """用于实现文件夹的自然排序"""
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', str(s))]

def main():
    # 创建实例
    extractor = MangaPanelExtractor(white_threshold=230, min_line_thickness=3)
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
   # 获取所有JPG文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg') or f.lower().endswith('.png')]
    
    print(f"\n找到 {len(image_files)} 个图片文件")
    
    if not image_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到图片文件！")
        return
    
    #进行自然排序
    image_files.sort(key=natural_sort_key)
    
    # 处理每个图片文件
    for image_file in image_files:
        print(f"\n处理文件: {image_file}")
        input_path = os.path.join(input_dir, image_file)
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(image_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {image_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {image_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 